import websocket
import json
import socket
import uuid
import threading
import time
import logging
import sys

import cv2
import numpy as np
from ultralytics import YOLO


# sys.path.append('/home/pi/TurboPi/')
# import HiwonderSDK.mecanum as mecanum
# import HiwonderSDK.Board as Board

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

turnOffFlag = False;
model = None


class RobustDeviceWebSocketClient:
    def __init__(self, server_url, username="robust_device"):
        self.server_url = server_url
        self.username = username
        self.device_ip = self.get_device_ip()
        self.mac_address = self.get_mac_address()
        self.ws = None
        self.is_connected = False
        self.reconnect_interval = 5  # 重连间隔（秒）
        self.max_reconnect_attempts = 10
        self.reconnect_attempts = 0
        # 加载YOLO模型
        global model
        logger.info("正在加载YOLO模型...")  # 使用logger代替print
        model = load_yolo_model()
        logger.info("模型加载完成！")  # 使用logger代替print

    def get_device_ip(self):
        """获取设备IP地址"""
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except Exception as e:
            logger.error(f"获取IP地址失败: {e}")
            return socket.gethostbyname(socket.gethostname())

    def get_mac_address(self):
        """获取MAC地址"""
        try:
            mac = ':'.join(['{:02x}'.format((uuid.getnode() >> elements) & 0xff)
                            for elements in range(0, 2 * 6, 2)][::-1])
            return mac
        except Exception as e:
            logger.error(f"获取MAC地址失败: {e}")
            return "unknown"

    def build_websocket_url(self):
        """构建WebSocket连接URL"""
        url = f"{self.server_url}/webSocket/{self.username}?ip={self.device_ip}&mac={self.mac_address}"
        return url

    def on_message(self, ws, message):
        """接收消息回调"""
        logger.info(f"收到服务器消息: {message}")
        try:
            self.handle_server_message(message)
        except Exception as e:
            logger.error(f"处理消息异常: {e}")

    def on_error(self, ws, error):
        """错误处理回调"""
        logger.error(f"WebSocket错误: {error}")
        self.is_connected = False

    def on_close(self, ws, close_status_code, close_msg):
        """连接关闭回调"""
        logger.info(f"WebSocket连接已关闭 - 状态码: {close_status_code}, 消息: {close_msg}")
        self.is_connected = False
        self.attempt_reconnect()

    def on_open(self, ws):
        """连接建立回调"""
        logger.info(f"WebSocket连接成功 - 设备IP: {self.device_ip}, MAC: {self.mac_address}")
        self.is_connected = True
        self.reconnect_attempts = 0

        # 发送设备上线信息
        self.send_device_info()

    def handle_server_message(self, message):
        """处理服务器消息"""
        try:
            data = json.loads(message)
            type_str = data.get("type")
            if type_str == "move":
                #车辆移动
                degree = data.get("degree")
                if degree is not None:
                    # chassis.set_velocity(50,degree,0)
                    time.sleep(0.1)
                    # chassis.set_velocity(0,0,0)
            
            elif type_str == "crosswise":
                #舵机横向移动
                crosswiseValue = data.get("crosswiseValue")
                if crosswiseValue is not None:
                    print("舵机横向移动")
                    #Board.setPWMServoPulse(2, crosswiseValue, 300)
                
            elif type_str == "vertical":
                #舵机竖向移动
                verticalValue = data.get("verticalValue")
                if verticalValue is not None:
                    print("舵机竖向移动")
                    # Board.setPWMServoPulse(1, verticalValue, 300)
            
            elif type_str == "reset":
                print("舵机竖向移动")
                #舵机竖向移动
                # Board.setPWMServoPulse(1, 1500, 300)
                # Board.setPWMServoPulse(2, 1500, 300)

            elif type_str == "color":
                #led灯颜色
                rgb = data.get("rgb")
                r = rgb.get("r")
                g = rgb.get("g")
                b = rgb.get("b")
                # Board.RGB.setPixelColor(0, Board.PixelColor(r, g, b))
                # Board.RGB.setPixelColor(1, Board.PixelColor(r, g, b))
                # Board.RGB.show()
                time.sleep(1)

            elif type_str == "discernOn":
                #开启智能识别
                global turnOffFlag
                turnOffFlag = True
                logger.info("已开启视频检测功能")
                time.sleep(1)

            elif type_str == "discernOff":
                # 关闭智能识别
                turnOffFlag = False
                logger.info("已关闭视频检测功能")
                time.sleep(1)

                
            
            # msg_type = data.get("type", "unknown")
            #
            # if msg_type == "ping":
            #     # 响应心跳
            #     self.send_pong()
            # elif msg_type == "command":
            #     # 处理命令
            #     self.execute_command(data.get("command"))
            # else:
            #     logger.info(f"未知消息类型: {msg_type}")

        except json.JSONDecodeError:
            logger.warning(f"无法解析JSON消息: {message}")

    def send_device_info(self):
        """发送设备信息"""
        device_info = {
            "type": "device_info",
            "ip": self.device_ip,
            "mac": self.mac_address,
            "status": "online",
            "timestamp": int(time.time())
        }
        self.send_message(json.dumps(device_info))

    def send_pong(self):
        """发送心跳响应"""
        pong_data = {
            "type": "pong",
            "timestamp": int(time.time())
        }
        self.send_message(json.dumps(pong_data))

    def execute_command(self, command):
        """执行服务器命令"""
        logger.info(f"执行命令: {command}")
        # 在这里实现具体的命令执行逻辑

        # 发送命令执行结果
        result = {
            "type": "command_result",
            "command": command,
            "result": "success",
            "timestamp": int(time.time())
        }
        self.send_message(json.dumps(result))

    def send_message(self, message):
        """发送消息"""
        if self.ws and self.is_connected:
            try:
                self.ws.send(message)
                logger.debug(f"发送消息: {message}")
            except Exception as e:
                logger.error(f"发送消息失败: {e}")
                self.is_connected = False
        else:
            logger.warning("WebSocket未连接，无法发送消息")

    def connect(self):
        """建立WebSocket连接"""
        ws_url = self.build_websocket_url()
        logger.info(f"连接到: {ws_url}")

        self.ws = websocket.WebSocketApp(
            ws_url,
            on_open=self.on_open,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close
        )

        self.ws.run_forever()

    def attempt_reconnect(self):
        """尝试重连"""
        if self.reconnect_attempts < self.max_reconnect_attempts:
            self.reconnect_attempts += 1
            logger.info(f"尝试重连 ({self.reconnect_attempts}/{self.max_reconnect_attempts})...")
            time.sleep(self.reconnect_interval)

            # 在新线程中重连
            reconnect_thread = threading.Thread(target=self.connect)
            reconnect_thread.daemon = True
            reconnect_thread.start()
        else:
            logger.error("达到最大重连次数，停止重连")



    def start(self):
        """启动客户端"""
        try:
            connect_thread = threading.Thread(target=self.connect)
            connect_thread.daemon = True
            connect_thread.start()

            # 主线程保持运行
            while True:
                time.sleep(1)
                # 可以在这里添加其他业务逻辑

        except KeyboardInterrupt:
            logger.info("程序被用户中断")
            if self.ws:
                self.ws.close()



def load_yolo_model():
    global model
    """加载YOLOv5模型"""
    # 加载预训练的YOLOv5模型（nano版本，体积小、速度快）
    # model = torch.hub.load('ultralytics/yolov5', 'yolov5n', pretrained=True)
    model = YOLO("../model/yolo11n.pt", verbose=False)  # 添加verbose=False禁用模型初始化输出

    # 注释掉model.info()以避免输出模型信息
    # model.info()
    # 设置识别置信度阈值（过滤低置信度结果）
    model.conf = 0.4  # 置信度>0.4才会被识别
    return model

def detect_objects(image):
    global turnOffFlag,model
    """
    识别图像中的目标并绘制边界框（适配YOLO11）

    Args:
        image: numpy.ndarray - 输入的BGR格式图像（OpenCV读取）
        model: ultralytics.YOLO - 加载好的YOLO11模型实例
    Returns:
        numpy.ndarray - 绘制了检测框、类别标签和置信度的结果图像
    """
    # 输入有效性校验
    if not isinstance(image, np.ndarray) or len(image.shape) != 3:
        raise ValueError("输入必须是OpenCV格式的BGR彩色图像（numpy.ndarray）")

    # 复制原图避免修改输入
    result_img = image.copy()


    # 判断是否开启了图像检测的方法
    if turnOffFlag:
        # YOLO11推理检测，添加verbose=False禁用推理输出
        results = model(image, verbose=False)

        # 处理检测结果（单张图像场景）
        if not results or len(results) == 0:
            return result_img

        # 获取第一张图像的检测结果
        result = results[0]

        # 检查是否存在检测框
        if not hasattr(result, 'boxes') or result.boxes is None:
            return result_img

        # 提取检测框数据（转为CPU numpy数组，避免GPU内存占用）
        boxes_data = result.boxes.data.cpu().numpy()

        # 遍历所有检测框
        for det in boxes_data:
            # YOLO11检测框格式：[x1, y1, x2, y2, conf, cls]（无name字段，需通过模型映射）
            x1, y1, x2, y2, conf, cls_id = det[:6]

            # 坐标转为整数（像素位置必须为int）
            x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])

            # 通过模型类别字典获取类别名称（YOLO11标准方式）
            cls_name = model.names[int(cls_id)] if hasattr(model, 'names') else f"Class_{int(cls_id)}"

            # 绘制检测框（绿色，线宽2）
            cv2.rectangle(result_img, (x1, y1), (x2, y2), (176, 207, 176), 2)

            # 标签位置优化（避免超出图像边界）
            label_y = y1 - 10 if y1 - 10 > 10 else y1 + 20

            # 绘制标签背景（提升可读性）
            label_text = f"{cls_name} {conf:.2f}"
            (text_w, text_h), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
            cv2.rectangle(result_img, (x1, label_y - text_h - 5),
                          (x1 + text_w, label_y + 5), (0, 255, 0), -1)

            # 绘制标签文字（白色字体对比更清晰）
            cv2.putText(result_img, label_text, (x1, label_y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)


    return result_img

# 使用示例
# if __name__ == "__main__":
#     # chassis = mecanum.MecanumChassis()
#     server_url = "ws://192.168.3.54:3570/ruoyi-vue-service"
#     client = RobustDeviceWebSocketClient(server_url, username="car")
#     client.start()
#     # chassis.set_velocity(0,0,0)
